"""
MCP Client 实现
这个模块实现了一个智能对话客户端，可以连接到 MCP 服务器并使用通义千问大模型处理用户查询
"""

import asyncio
from typing import Optional, Dict, Any
from contextlib import AsyncExitStack
import logging
import sys
import json
import signal
from asyncio import TimeoutError

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

import dashscope
from dashscope import Generation
from dotenv import load_dotenv
import os

# 配置日志系统
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 加载环境变量
load_dotenv()

class MCPClientError(Exception):
    """MCP 客户端自定义异常"""
    pass

class MCPClient:
    """
    MCP 客户端类
    负责管理与服务器的连接、处理用户查询、协调工具调用等核心功能
    """
    
    def __init__(self):
        """
        初始化 MCP 客户端
        设置必要的属性并验证环境变量
        """
        # 初始化会话和客户端对象
        self.session: Optional[ClientSession] = None
        self.exit_stack = AsyncExitStack()
        self.available_tools = {}
        
        # 获取并验证 API 密钥
        self.api_key = os.getenv('DASHSCOPE_API_KEY')
        if not self.api_key:
            raise ValueError("DASHSCOPE_API_KEY not found in environment variables")
        logger.info("MCPClient initialized with API key")
        
        # 设置 DashScope API Key
        dashscope.api_key = self.api_key

    async def connect_to_server(self, server_script_path: str):
        """
        连接到 MCP 服务器
        
        参数:
            server_script_path: 服务器脚本路径 (.py 或 .js)
        """
        logger.info(f"Attempting to connect to server with script: {server_script_path}")
        
        # 检查脚本类型
        if not os.path.exists(server_script_path):
            raise FileNotFoundError(f"Server script not found: {server_script_path}")
            
        is_python = server_script_path.endswith('.py')
        is_js = server_script_path.endswith('.js')
        if not (is_python or is_js):
            raise ValueError("Server script must be a .py or .js file")
            
        # 根据脚本类型选择命令
        if is_python:
            # 使用 Python 3.10 的完整路径
            command = "/opt/homebrew/Cellar/python@3.10/3.10.17/bin/python3.10"
            if not os.path.exists(command):
                command = "python3"  # 如果找不到特定版本，使用 python3
        else:
            command = "node"
            
        server_params = StdioServerParameters(
            command=command,
            args=[server_script_path],
            env=None
        )
        
        # 建立服务器连接
        logger.info(f"Establishing connection to server using {command}...")
        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        self.stdio, self.write = stdio_transport
        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))
        
        # 等待服务器启动
        await asyncio.sleep(1)
        
        # 初始化连接
        await self.session.initialize()
        logger.info("Server connection initialized successfully")
        
        # 获取并缓存可用工具列表
        try:
            response = await self.session.list_tools()
            self.available_tools = {
                tool.name: {
                    "description": tool.description,
                    "parameters": tool.inputSchema
                }
                for tool in response.tools
            }
            tool_names = list(self.available_tools.keys())
            logger.info(f"Available tools: {tool_names}")
            print("\nConnected to server with tools:", tool_names)
            print("\nServer is ready for queries!")
        except Exception as e:
            logger.error(f"Failed to get tool list: {str(e)}")
            raise

    async def _execute_tool_call(self, tool_call) -> str:
        """
        执行工具调用
        
        参数:
            tool_call: 工具调用信息
            
        返回:
            str: 工具调用结果
        """
        try:
            # 获取工具名称和参数
            if isinstance(tool_call, dict):
                if 'function' in tool_call:
                    tool_name = tool_call['function']['name']
                    parameters = tool_call['function']['arguments']
                else:
                    tool_name = tool_call.get('name')
                    parameters = tool_call.get('arguments', '{}')
            else:
                tool_name = tool_call.function.name
                parameters = tool_call.function.arguments
            
            # 解析参数
            try:
                if isinstance(parameters, str):
                    parameters = json.loads(parameters)
                elif isinstance(parameters, dict):
                    parameters = parameters
                else:
                    parameters = {"text": str(parameters)}
            except json.JSONDecodeError:
                parameters = {"text": parameters}
            
            # 记录工具调用信息
            logger.info(f"Calling tool: {tool_name} with args: {parameters}")
            
            # 执行工具调用
            result = await self.session.call_tool(tool_name, parameters)
            
            # 处理返回结果
            if isinstance(result.content, str):
                try:
                    content = json.loads(result.content)
                    return json.dumps(content, ensure_ascii=False, indent=2)
                except json.JSONDecodeError:
                    return result.content
            return str(result.content)
            
        except Exception as e:
            error_msg = f"Tool call error: {str(e)}"
            logger.error(error_msg)
            return error_msg

    async def process_query(self, query: str) -> str:
        """
        处理用户查询
        
        使用通义千问模型处理查询，支持工具调用
        
        参数:
            query: 用户输入的查询字符串
            
        返回:
            str: 处理结果
        """
        logger.info(f"Processing query: {query}")
        
        if not self.session:
            raise MCPClientError("No active server connection")
            
        # 构建对话消息
        messages = [
            {
                "role": "user",
                "content": query
            }
        ]

        try:
            # 准备工具列表
            available_tools = [
                {
                    "type": "function",
                    "function": {
                        "name": name,
                        "description": info["description"],
                        "parameters": info["parameters"]
                    }
                }
                for name, info in self.available_tools.items()
            ]

            # 记录请求参数
            request_params = {
                "model": 'qwen-max',
                "messages": messages,
                "tools": available_tools,
                "result_format": 'message'
            }
            logger.info(f"Calling Qwen API with params: {json.dumps(request_params, ensure_ascii=False)}")

            # 调用通义千问模型
            try:
                response = Generation.call(
                    model='qwen-max',
                    messages=messages,
                    tools=available_tools,
                    result_format='message',
                    api_key=self.api_key
                )
                
                # 记录响应状态
                logger.info(f"API response status: {response.status_code}")
                if response.status_code != 200:
                    logger.error(f"API error response: {response.message if hasattr(response, 'message') else 'Unknown error'}")
                    return f"API调用失败 (状态码: {response.status_code}): {response.message if hasattr(response, 'message') else '未知错误'}"
                
            except Exception as api_error:
                logger.error(f"API call error: {str(api_error)}")
                return f"API调用出错: {str(api_error)}"

            # 处理模型响应
            final_text = []
            
            # 获取响应内容
            message = response.output.choices[0].message
            logger.info(f"Model response message: {message}")
            
            # 将消息转换为字典格式
            if isinstance(message, dict):
                message_dict = message
            else:
                message_dict = message.__dict__ if hasattr(message, '__dict__') else {}
            
            # 获取内容
            content = message_dict.get('content', '')
            if content:
                final_text.append(content)
            
            # 处理工具调用请求
            function_call = message_dict.get('function_call')
            tool_calls = message_dict.get('tool_calls')
            
            if function_call:
                # 处理单个工具调用
                tool_name = function_call.get('name')
                tool_args = function_call.get('arguments')
                logger.info(f"Processing single function call: {tool_name}")
                
                result = await self._execute_tool_call({
                    "type": "function",
                    "function": {
                        "name": tool_name,
                        "arguments": tool_args
                    }
                })
                final_text.append(f"\nTool result: {result}")
                
                # 继续对话
                messages.extend([
                    {
                        "role": "assistant",
                        "content": content,
                        "function_call": {
                            "name": tool_name,
                            "arguments": tool_args
                        }
                    },
                    {
                        "role": "function",
                        "name": tool_name,
                        "content": result
                    }
                ])
                
            elif tool_calls:
                # 处理多个工具调用
                for tool_call in tool_calls:
                    result = await self._execute_tool_call(tool_call)
                    final_text.append(f"\nTool result: {result}")
                    
                    # 继续对话
                    messages.extend([
                        {
                            "role": "assistant",
                            "content": content,
                            "tool_calls": [tool_call]
                        },
                        {
                            "role": "tool",
                            "content": result,
                            "tool_call_id": tool_call.get('id')
                        }
                    ])

            # 获取最终响应
            if len(messages) > 1:
                try:
                    final_response = Generation.call(
                        model='qwen-max',
                        messages=messages,
                        tools=available_tools,
                        result_format='message',
                        api_key=self.api_key
                    )
                    
                    if final_response.status_code == 200:
                        final_message = final_response.output.choices[0].message
                        if isinstance(final_message, dict):
                            final_content = final_message.get('content', '')
                        else:
                            final_content = final_message.content if hasattr(final_message, 'content') else ''
                        
                        if final_content:
                            final_text.append(final_content)
                    else:
                        logger.error(f"Final API call failed: {final_response.message if hasattr(final_response, 'message') else 'Unknown error'}")
                except Exception as final_error:
                    logger.error(f"Final API call error: {str(final_error)}")
                    final_text.append(f"\nError in final response: {str(final_error)}")

            return "\n".join(final_text) if final_text else "抱歉，我没有得到有效的响应。"
            
        except Exception as e:
            logger.error(f"Error processing query: {str(e)}")
            return f"处理查询时出错: {str(e)}"

    async def chat_loop(self):
        """
        运行交互式对话循环
        
        持续接收用户输入并处理查询，直到用户输入 'quit'
        """
        print("\nMCP Client Started!")
        print("Type your queries or 'quit' to exit.")
        
        while True:
            try:
                query = input("\nQuery: ").strip()
                
                if query.lower() == 'quit':
                    break
                    
                response = await self.process_query(query)
                print("\nResponse:", response)
                    
            except Exception as e:
                print(f"\nError: {str(e)}")
    
    async def cleanup(self):
        """
        清理资源
        
        关闭所有打开的连接和资源
        """
        if self.exit_stack:
            await self.exit_stack.aclose()

async def main():
    """
    主函数
    
    处理命令行参数并运行客户端
    """
    if len(sys.argv) < 2:
        print("Usage: python client.py <path_to_server_script>")
        sys.exit(1)
        
    client = MCPClient()
    try:
        # 连接服务器并运行对话循环
        await client.connect_to_server(sys.argv[1])
        await client.chat_loop()
    except FileNotFoundError as e:
        logger.error(f"Server script not found: {e}")
        print(f"\nError: {e}")
    except TimeoutError as e:
        logger.error(f"Connection timeout: {e}")
        print(f"\nError: Server connection timed out. Please check if:")
        print("1. The server script path is correct")
        print("2. The server script has no errors")
        print("3. The server script is properly implementing the MCP protocol")
    except Exception as e:
        logger.error(f"Fatal error: {str(e)}")
        print(f"\nFatal error: {str(e)}")
    finally:
        await client.cleanup()

if __name__ == "__main__":
    import sys
    asyncio.run(main())